import uuid
from datetime import datetime

from sqlmodel import SQLModel, Field, Relationship
from app.models.user_models import User


class TrainingSessionBase(SQLModel):
    topic: str
    training_type: str  # 培训类型，如"立论训练"、"质询训练"等


class TrainingSessionCreate(TrainingSessionBase):
    pass


class TrainingSessionUpdate(SQLModel):
    topic: str | None = None
    training_type: str | None = None


class TrainingSession(TrainingSessionBase, table=True):
    __tablename__ = "training_sessions"
    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    user_id: uuid.UUID = Field(foreign_key="users.id", nullable=False)
    created_at: datetime = Field(default_factory=datetime.now)

    user: User | None = Relationship(back_populates="training_sessions")
    turns: list["TrainingTurn"] = Relationship(back_populates="session")


class TrainingSessionPublic(TrainingSessionBase):
    id: uuid.UUID
    user_id: uuid.UUID
    created_at: datetime


class TrainingSessionsPublic(SQLModel):
    data: list[TrainingSessionPublic]
    count: int


class TrainingTurn(SQLModel, table=True):
    __tablename__ = "training_turns"
    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    session_id: uuid.UUID = Field(foreign_key="training_sessions.id", nullable=False)
    speaker: str  # 'user' or 'ai'
    content: str
    turn_index: int
    timestamp: datetime = Field(default_factory=datetime.now)

    session: TrainingSession = Relationship(back_populates="turns")


class TrainingTurnCreate(SQLModel):
    session_id: uuid.UUID
    speaker: str
    content: str
    turn_index: int


class TrainingTurnPublic(SQLModel):
    id: uuid.UUID
    session_id: uuid.UUID
    speaker: str
    content: str
    turn_index: int
    timestamp: datetime


class TrainingHistoryResponse(SQLModel):
    """用于返回培训历史记录的响应模型"""
    data: list[TrainingSessionPublic]
    count: int


class TrainingTurnsResponse(SQLModel):
    """用于返回培训对话记录的响应模型"""
    data: list[TrainingTurnPublic]
    count: int